Repository navigation
feat(libtorch): piecewise-linear table lookup composed from float32 primitives - #34
NicolasRouquette wants to merge 2 commits into
Conversation
Brings in the upstream PR branch (lean-dojo#34) at 558a3af, one commit on upstream main b062b9a: floor in operations.h, Buffer.gatherAt with its scatterAddAt VJP, Tape.floor, Tape.gatherAt, and Tape.tableLookup / Tape.layeredTableLookup composed from float32 primitives so the forward matches KernelSpec.tableLookupSpec bit for bit. Successor of the closed lean-dojo#24 (TexTable). The texture mechanism and the hardware filter mode do not carry over to this backend.
|
Hey Nicolas, thanks for the interpolation work. I found a bounds issue in table_positions: the upper clamp bound is converted to float32 before the integer conversion. With size = 16777220 and float32 positions [inf, 16777220], the same ATen nan_to_num/clamp/to(int64) sequence produces indices [16777220, 16777220], but the last valid index is 16777219. tableLookup limits its table size, but the exposed gatherAt/scatterAddAt primitives do not have that limit. Could you clamp again in the integer domain to [0, size - 1] and add a regression above 2^24, including infinity? Holding off on merging until this is fixed. |
…rimitives Add `floor` to the shared operation list, a gather at device-resident positions with its scatter-add VJP, and `Tape.tableLookup` / `Tape.layeredTableLookup`: clamp-addressed linear interpolation into a tabulated function, composed in Lean from `clamp`, `floor`, `sub`, `mul`, `add`, and the new gather so that every operation rounds once and the result matches `KernelSpec.tableLookupSpec` bit for bit. The backward gives the coordinate its segment's slope, zero where the clamp is active, and the table the two interpolation weights. Tests compare `floor` with `Float32.floor`, `gatherAt` with the CPU tape's `indexSelect`, the lookups with a `Float32` reference at boundary, out-of-range, integral, and NaN coordinates, and the forward and both gradients with the same function recorded on the CPU tape. A repeated `scatterAddAt` must return the same bits under the deterministic control.
`table_positions` clamped float32 positions to `[0, size - 1]` and converted the result to int64, but above 2^24 the bound `size - 1` has no float32 representation and rounds up to `size` (16777219 reads as 16777220), so an infinite or oversize position selected one past the end. `gatherAt` and `scatterAddAt` are not bounded by the 2^24 limit that `Tape.tableLookup` enforces for exact positions. The index is now clamped again as int64, where every table index is exact, and the docstrings say so. The regression gathers from and scatter-adds into a table of 16777220 entries at infinity, at 16777220, at the unrepresentable 16777219, at 1e9, and at NaN and negative positions; without the integer clamp it dies in `index_select` with a device-side assert.
558a3af to
11d1e3a
Compare
|
Confirmed and fixed; thank you for the exact numbers. const auto clamped = at::clamp(at::nan_to_num(positions, 0.0, last, 0.0), 0.0, last);
return at::clamp(clamped.to(at::kLong), int64_t{0}, static_cast<int64_t>(size) - 1);The regression ( I also ran the new test against the bridge without the integer clamp, to make sure it is the test (ATen's own Rebased onto current
Checks on it: RTX A4500, pip torch 2.11.0+cu128, CUDA 12.8 toolkit, Lean v4.35.0: The PR description is updated to match. |
What
A tabulated function (a calibration curve, an inverted sensor response, an activation lookup)
evaluated at grid coordinates by clamp-addressed linear interpolation: the uniform-grid analogue
of
np.interp, as a CUDA tape node with gradients for both the table and the coordinates.Tape.tableLookup {width} tableId coordsIdfor a rank-one table ofwidthsamples, andTape.layeredTableLookup {layers width} tableId coordsId layerIdfor alayers × widthtablewhose layer is selected per element from a third tape value. The result has the coordinates'
shape. With
i = ⌊u⌋andw = u - iafter clampinguto[0, width - 1], the value isv₀ + (v₁ - v₀) · wwithv₀ = table[i]andv₁ = table[min (i + 1) (width - 1)].floorin the sharedoperations.hlist (so the unavailable build gets its twin for free), andBuffer.gatherAtwith its VJP
Buffer.scatterAddAt, a gather at positions that live on the device as a float32buffer. Both also get tape nodes (
Tape.floor,Tape.gatherAt).KernelSpec.tableLookupSpec, theExecFloat.Binary 8 23reference the forward matches.Relation to #24
#24 put the same function behind a layered
cudaArrayand a texture object, which is the kind ofcustom CUDA this backend no longer has. The function survives the move; the mechanism does not.
What this PR keeps from #24 is the contract that mattered there: the forward is bit-identical to
an executable
Float32reference, because every operation it is composed from rounds exactlyonce (each ATen elementwise kernel is one rounding, and nothing contracts across kernel launches).
What it adds is the backward #24 deferred: the coordinate receives the slope of its segment, zero
where the clamp is active (the same open-interval rule as
clamp), and the table receives the twointerpolation weights scattered to the samples they multiplied. The 9-bit hardware filter mode
is not reproducible on ATen and is not offered.
Design notes
gatherAttakes positions as a float32 buffer, which is what the tapecomputes: a NaN position reads as
0, every position is clamped to[0, size - 1]and thentruncated, so a non-integral position floors and no selection leaves the table without a
device-side assertion. The clamp is applied to the float32 position and again to the integer
index: above
2^24the boundsize - 1has no float32 representation and rounds up tosize(16777219 reads as 16777220), so the float32 clamp alone would let an infinite or oversize
position select one past the end. The tape operation rejects empty tables and tables above
2^24samples, which is where float32 positions stop being exact; the two primitives have nosuch limit and stay in bounds without it.
tableSegmentscomputes the clamped coordinate, its floor, theweight, and the neighbouring position with
clamp,floor,sub,add; all of these areexact in float32 below
2^24. Only the finalsub,mul,addround, in that order, which isthe order
tableLookupSpecstates.cleanuplist for thebackward, so backward does not recompute them, and the eager session releases them with the
tape. Every other intermediate is released through the value that consumed it.
is a correctness and composability change, not a performance claim, and no timing is reported.
scatterAddAtaccumulates repeated positions throughindex_add_,whose order is the SDK's. A repeatability check under the deterministic control is included
with the other deterministic reductions; the gradient comparison uses dyadic values, where the
order does not change the sum.
Tests
NN/Tests/Runtime/Cuda/TableLookup.lean, wired into the CUDA coverage suite:flooragrees withFloat32.floorbit for bit, including signed zeros,16777217, and±1e8; its tape node reports a zero gradient.gatherAtagrees with the CPU tape'sindexSelectexactly, forward and table gradient,including a repeated position; and the documented addressing holds for negative, past-the-end,
non-integral, and NaN positions.
gatherAtandscatterAddAtstay in bounds on a table of 16777220 entries (above2^24): at+inf, at16777220.0, at16777219.0(which narrows to16777220.0), at1e9, at-inf,at NaN and at a negative position, each selecting the entry the documented addressing names.
Without the integer clamp this test dies in
index_selectwith a device-side assert.Both lookups match the
Float32reference bit for bit over interior, boundary, out-of-range,integral, signed-zero, and NaN coordinates, with layer selectors past the end and non-integral.
The forward and both gradients agree with the same function recorded on the CPU tape from
existing nodes (
clamp,sub,indexSelecttwice,sub,mul,add, with the segmentindices computed on the host by the same recipe).
Edges: a one-sample table (every coordinate reads it, no coordinate gradient, the table
gradient is the seed sum), zero coordinates, and the two rejections.
scatterAddAtaccumulates with atomics, so the same file checks that a repeated call returnsthe same bits under the deterministic control (and restores the previous settings).
Checks run
On
11d1e3a6(two commits onmain72754452). RTX A4500 (sm_86, driver 580.126), pip torch2.11.0+cu128 as the SDK, CUDA 12.8 toolkit for SDK discovery and NVRTC, Lean v4.35.0:
Negative control on the same tip, the test kept and the bridge rebuilt without the integer clamp:
Not run: the sanitizer harness and the elementwise C++ harness.